#!/usr/bin/env python3
"""
油田态势对象数据导入脚本
从 Excel 文件读取油田数据并导入到 PostgreSQL 数据库
"""

import sys
import os
import logging
import asyncio
from pathlib import Path
from datetime import datetime
from typing import Dict, Any, List

# 尝试导入 pandas 和 openpyxl
try:
    import pandas as pd
    HAS_PANDAS = True
except ImportError:
    HAS_PANDAS = False
    logging.warning("pandas 未安装，将尝试使用其他方式读取Excel文件")

try:
    import openpyxl
    HAS_OPENPYXL = True
except ImportError:
    HAS_OPENPYXL = False
    logging.warning("openpyxl 未安装，将尝试使用其他方式读取Excel文件")

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# 直接导入所需模块，避免触发 src.__init__ 中的图数据库初始化
import src.situation.database
import src.situation.models
import src.situation.service
import src.situation.schemas

from src.situation.database import get_db_session, engine, Base
from src.situation.models import SituationObject
from src.situation.service import SituationService
from src.situation.schemas import SituationObjectCreate, DeploymentStatus, ObjectType


# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('logs/oilfield_import.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


def read_excel_file(file_path: str) -> List[Dict[str, Any]]:
    """
    读取 Excel 文件并返回数据列表

    Args:
        file_path: Excel 文件路径

    Returns:
        包含油田数据的字典列表
    """
    try:
        if HAS_PANDAS and HAS_OPENPYXL:
            # 使用 pandas 读取 Excel 文件
            df = pd.read_excel(file_path)
            logger.info(f"使用 pandas 成功读取 Excel 文件，共 {len(df)} 行数据")

            # 清理列名中的不可见字符
            df.columns = [col.strip().replace('\xa0', '') for col in df.columns]

            # 转换为字典列表
            data = df.to_dict('records')

            # 记录列名
            logger.info(f"Excel 文件包含的列: {list(df.columns)}")

            return data
        else:
            # 使用 openpyxl 直接读取
            if not HAS_OPENPYXL:
                raise ImportError("需要安装 pandas 和 openpyxl 来读取 Excel 文件")

            from openpyxl import load_workbook

            wb = load_workbook(filename=file_path, read_only=True)
            sheet = wb.active

            # 读取标题行并清理列名
            headers = [cell.value.strip().replace('\xa0', '') if cell.value else '' for cell in next(sheet.iter_rows(min_row=1, max_row=1))]
            logger.info(f"Excel 文件包含的列: {headers}")

            # 读取数据行
            data = []
            for row in sheet.iter_rows(min_row=2, values_only=True):
                if any(row):  # 跳过空行
                    row_dict = {headers[i]: row[i] for i in range(len(headers)) if i < len(row)}
                    data.append(row_dict)

            wb.close()
            logger.info(f"使用 openpyxl 成功读取 Excel 文件，共 {len(data)} 行数据")
            return data

    except Exception as e:
        logger.error(f"读取 Excel 文件失败: {e}")
        if not HAS_PANDAS or not HAS_OPENPYXL:
            logger.error("请安装必要的依赖: uv add pandas openpyxl")
        raise


def parse_coordinates(coord_str: str) -> Dict[str, float]:
    """
    解析经纬度字符串

    Args:
        coord_str: 经纬度字符串，格式如 "25.43°N 49.62°E" 或 "约 28.3°N, 48.8°E"

    Returns:
        包含 latitude 和 longitude 的字典
    """
    if not coord_str:
        return {}

    try:
        import re
        # 移除"约"等前缀
        coord_str = coord_str.replace('约', '').strip()

        # 提取纬度和经度
        lat_match = re.search(r'([\d.]+)\s*°?\s*N', coord_str, re.IGNORECASE)
        lon_match = re.search(r'([\d.]+)\s*°?\s*E', coord_str, re.IGNORECASE)

        result = {}
        if lat_match:
            result['latitude'] = float(lat_match.group(1))
        if lon_match:
            result['longitude'] = float(lon_match.group(1))

        return result
    except Exception as e:
        logger.warning(f"解析经纬度失败: {coord_str}, 错误: {e}")
        return {}


def convert_to_situation_object(oilfield_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    将油田数据转换为态势对象格式

    Args:
        oilfield_data: 油田原始数据

    Returns:
        符合态势对象模型的数据字典
    """
    try:
        # 基础信息映射
        name = oilfield_data.get('名称') or oilfield_data.get('name')
        country = oilfield_data.get('国家') or oilfield_data.get('country')
        coordinates = oilfield_data.get('经纬度') or oilfield_data.get('coordinates')

        # 使用名称作为编码的一部分
        if name:
            # 生成唯一编码
            import re
            import random
            # 提取英文部分
            english_part = re.search(r'\(([^)]+)\)', name)
            if english_part:
                code_base = english_part.group(1).replace(' ', '_').upper()
            else:
                code_base = f"OILFIELD_{datetime.now().strftime('%Y%m%d%H%M%S')}"

            code = f"{code_base}_{random.randint(1000, 9999)}"
        else:
            name = f"未知油田-{datetime.now().strftime('%Y%m%d%H%M%S')}"
            code = f"OILFIELD-{datetime.now().strftime('%Y%m%d%H%M%S')}"

        # 解析坐标
        location_info = {}
        if coordinates:
            coords = parse_coordinates(coordinates)
            if coords:
                location_info.update(coords)
            location_info['description'] = coordinates

        # 油田特定数据
        discovery_time = oilfield_data.get('发现时间') or oilfield_data.get('discovery_time') or ''
        exploitation_time = oilfield_data.get('开采时间') or oilfield_data.get('exploitation_time') or ''
        total_reserves = oilfield_data.get('总储量') or oilfield_data.get('total_reserves') or ''
        annual_output = oilfield_data.get('年产量') or oilfield_data.get('annual_output') or ''
        oilfield_area = oilfield_data.get('油田面积') or oilfield_data.get('oilfield_area') or ''
        main_production_layer = oilfield_data.get('主产油层') or oilfield_data.get('main_production_layer') or ''

        # 构建态势对象数据
        situation_data = {
            'object_type': ObjectType.INFRASTRUCTURE,  # 使用枚举类型
            'name': name,
            'code': code,
            'organization': country,  # 将国家映射到所属单位
            'deployment_status': DeploymentStatus.DEPLOYED,  # 使用枚举

            # 油田特定数据存储在 type_specific_data 中
            'type_specific_data': {
                'infrastructure_category': '油田',  # 设施类别

                # 油田专属字段
                'discovery_time': discovery_time,
                'exploitation_time': exploitation_time,
                'total_reserves': total_reserves,
                'annual_output': annual_output,
                'oilfield_area': oilfield_area,
                'main_production_layer': main_production_layer,
            },

            # 适用场景
            'applicable_scenarios': ['能源开采', '石油生产', '工业生产'],

            # 标签
            'tags': ['油田', '基础设施', '能源'],

            # 其他元数据
            'extra_metadata': {
                'data_source': 'excel_import',
                'import_time': datetime.now().isoformat(),
                'country': country
            },

            # 审计信息
            'created_by': 'oilfield_import_script',
            'updated_by': 'oilfield_import_script'
        }

        # 添加位置信息
        if location_info:
            situation_data['location'] = location_info

        logger.debug(f"成功转换油田数据: {name} (编码: {code})")
        return situation_data

    except Exception as e:
        logger.error(f"转换油田数据失败: {e}")
        logger.error(f"原始数据: {oilfield_data}")
        raise


async def import_oilfield_data(excel_file_path: str, batch_size: int = 10):
    """
    导入油田数据到数据库

    Args:
        excel_file_path: Excel 文件路径
        batch_size: 批量处理大小
    """
    logger.info(f"=" * 80)
    logger.info(f"开始导入油田数据")
    logger.info(f"文件路径: {excel_file_path}")
    logger.info(f"=" * 80)

    try:
        # 读取 Excel 文件
        oilfield_data_list = read_excel_file(excel_file_path)
        logger.info(f"成功读取 Excel 文件，共 {len(oilfield_data_list)} 条油田数据")

        # 首先初始化数据库（确保表结构正确）
        logger.info("正在初始化数据库...")
        from src.situation.database import init_db
        await init_db()
        logger.info("数据库初始化完成")

        # 创建数据库会话
        async for db in get_db_session():
            success_count = 0
            error_count = 0
            skipped_count = 0
            error_details = []

            for i, oilfield_data in enumerate(oilfield_data_list, 1):
                try:
                    logger.info(f"\n{'=' * 60}")
                    logger.info(f"处理第 {i}/{len(oilfield_data_list)} 条数据")

                    # 转换为态势对象格式
                    situation_data = convert_to_situation_object(oilfield_data)

                    logger.info(f"油田名称: {situation_data['name']}")
                    logger.info(f"油田编码: {situation_data.get('code', '无')}")
                    logger.info(f"所属国家: {situation_data.get('organization', '无')}")

                    # 创建态势对象
                    obj_create = SituationObjectCreate(**situation_data)

                    # 创建态势对象 (SituationService.create_object 已经内置了编码重复检查)
                    created_object = await SituationService.create_object(
                        db, obj_create, user="oilfield_import_script"
                    )
                    logger.info(f"✓ 成功创建态势对象: {created_object.name} (ID: {created_object.id})")
                    success_count += 1

                except ValueError as e:
                    error_msg = str(e)
                    if "已存在" in error_msg or "重复" in error_msg:
                        logger.warning(f"✗ 数据已存在，跳过: {error_msg}")
                        skipped_count += 1
                        error_details.append({
                            'index': i,
                            'name': oilfield_data.get('名称', '未知'),
                            'reason': '数据已存在',
                            'error': error_msg
                        })
                    else:
                        logger.error(f"✗ 数据验证错误: {error_msg}")
                        error_count += 1
                        error_details.append({
                            'index': i,
                            'name': oilfield_data.get('名称', '未知'),
                            'reason': '数据验证失败',
                            'error': error_msg
                        })
                except Exception as e:
                    logger.error(f"✗ 处理第 {i} 条数据失败: {e}")
                    error_count += 1
                    error_details.append({
                        'index': i,
                        'name': oilfield_data.get('名称', '未知'),
                        'reason': '处理失败',
                        'error': str(e)
                    })
                    continue

            # 输出导入摘要
            logger.info(f"\n{'=' * 80}")
            logger.info("导入完成！")
            logger.info(f"{'=' * 80}")
            logger.info(f"总计: {len(oilfield_data_list)} 条")
            logger.info(f"成功: {success_count} 条")
            logger.info(f"跳过: {skipped_count} 条（已存在）")
            logger.info(f"失败: {error_count} 条")
            logger.info(f"{'=' * 80}")

            # 如果有错误，输出错误详情
            if error_details:
                logger.info("\n错误详情：")
                for detail in error_details:
                    logger.info(f"  第 {detail['index']} 条 - {detail['name']}: {detail['reason']}")
                    logger.debug(f"    错误信息: {detail['error']}")

            return {
                'total': len(oilfield_data_list),
                'success': success_count,
                'skipped': skipped_count,
                'error': error_count,
                'error_details': error_details
            }

    except Exception as e:
        logger.error(f"导入过程中发生严重错误: {e}")
        import traceback
        logger.error(traceback.format_exc())
        raise


async def main():
    """主函数"""
    # 默认文件路径
    default_file = "test/data-situation/oilfield.xlsx"

    # 检查命令行参数
    if len(sys.argv) > 1:
        excel_file = sys.argv[1]
    else:
        excel_file = default_file

    # 检查文件是否存在
    if not os.path.exists(excel_file):
        logger.error(f"文件不存在: {excel_file}")
        logger.info("用法: python scripts/import_oilfield_data.py [excel_file_path]")
        sys.exit(1)

    try:
        # 确保日志目录存在
        os.makedirs("logs", exist_ok=True)

        # 导入数据
        await import_oilfield_data(excel_file)

        logger.info("油田数据导入成功完成！")

    except Exception as e:
        logger.error(f"导入失败: {e}")
        sys.exit(1)


def run_main():
    """运行主函数的同步包装器"""
    asyncio.run(main())


if __name__ == "__main__":
    run_main()
